import torch, toml
import torchvision.models as models
import torch.nn as nn


def main(raw_path: str, prepared_path: str, config_path: str):
    # num_class = len(toml.load(config_path)["label"])
    num_class = 11

    # model = models.resnet34(weights=models.ResNet34_Weights.DEFAULT)
    model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
    # model = models.resnet101(weights=models.ResNet101_Weights.DEFAULT)
    torch.save(model, raw_path)

    # model.conv1 = nn.Conv2d(
    #     1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
    # )
    model.fc = nn.Linear(in_features=2048, out_features=num_class, bias=True)
    torch.save(model, prepared_path)


if __name__ == "__main__":
    main(
        "./model/raw/resnet50.pt",
        "./model/prepared/resnet50_3c.pt",
        "./data/config.toml",
    )

    # main(
    #     "./model/raw/resnet101.pt",
    #     "./model/prepared/resnet101.pt",
    #     "./data/config.toml",
    # )
